package pool

import (
	"github.com/panjf2000/ants/v2"
	"sync"
)

type Pool struct {
	*ants.Pool
	waitGroup *sync.WaitGroup
}

func (p *Pool) Submit(task func()) error {
	p.waitGroup.Add(1)
	return p.Pool.Submit(func() {
		defer p.waitGroup.Done()
		task()
	})
}

func (p *Pool) Wait() {
	p.waitGroup.Wait()
}

func NewPool(size int, options ...ants.Option) (*Pool, error) {
	pool, err := ants.NewPool(size, options...)
	if err != nil {
		return nil, err
	}
	return &Pool{
		Pool:      pool,
		waitGroup: new(sync.WaitGroup),
	}, nil
}
